package com.bruce.ai.alibaba.agentcard.discovery.circuit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Instant;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 简单的熔断器实现
 * 用于处理Agent调用失败的情况
 */
public class CircuitBreaker {
    
    private static final Logger logger = LoggerFactory.getLogger(CircuitBreaker.class);
    
    private final String name;
    private final int failureThreshold;
    private final long recoveryTimeoutMs;
    
    private final AtomicInteger failureCount = new AtomicInteger(0);
    private final AtomicLong lastFailureTime = new AtomicLong(0);
    private final AtomicReference<State> state = new AtomicReference<>(State.CLOSED);
    
    public CircuitBreaker(String name, int failureThreshold, long recoveryTimeoutMs) {
        this.name = name;
        this.failureThreshold = failureThreshold;
        this.recoveryTimeoutMs = recoveryTimeoutMs;
    }
    
    /**
     * 执行操作，如果熔断器开启则抛出异常
     */
    public <T> T execute(CircuitBreakerOperation<T> operation) throws Exception {
        if (!canExecute()) {
            throw new CircuitBreakerOpenException("Circuit breaker is OPEN for: " + name);
        }
        
        try {
            T result = operation.execute();
            onSuccess();
            return result;
        } catch (Exception e) {
            onFailure();
            throw e;
        }
    }
    
    /**
     * 检查是否可以执行操作
     */
    private boolean canExecute() {
        State currentState = state.get();
        
        switch (currentState) {
            case CLOSED:
                return true;
            case OPEN:
                if (shouldAttemptReset()) {
                    state.compareAndSet(State.OPEN, State.HALF_OPEN);
                    logger.info("Circuit breaker {} transitioned from OPEN to HALF_OPEN", name);
                    return true;
                }
                return false;
            case HALF_OPEN:
                return true;
            default:
                return false;
        }
    }
    
    /**
     * 成功时调用
     */
    private void onSuccess() {
        failureCount.set(0);
        State currentState = state.get();
        if (currentState == State.HALF_OPEN) {
            state.compareAndSet(State.HALF_OPEN, State.CLOSED);
            logger.info("Circuit breaker {} transitioned from HALF_OPEN to CLOSED", name);
        }
    }
    
    /**
     * 失败时调用
     */
    private void onFailure() {
        lastFailureTime.set(System.currentTimeMillis());
        int failures = failureCount.incrementAndGet();
        
        State currentState = state.get();
        if (currentState == State.HALF_OPEN) {
            state.compareAndSet(State.HALF_OPEN, State.OPEN);
            logger.warn("Circuit breaker {} transitioned from HALF_OPEN to OPEN", name);
        } else if (failures >= failureThreshold && currentState == State.CLOSED) {
            state.compareAndSet(State.CLOSED, State.OPEN);
            logger.warn("Circuit breaker {} transitioned from CLOSED to OPEN after {} failures", name, failures);
        }
    }
    
    /**
     * 检查是否应该尝试重置
     */
    private boolean shouldAttemptReset() {
        return System.currentTimeMillis() - lastFailureTime.get() >= recoveryTimeoutMs;
    }
    
    /**
     * 获取当前状态
     */
    public State getState() {
        return state.get();
    }
    
    /**
     * 获取失败次数
     */
    public int getFailureCount() {
        return failureCount.get();
    }
    
    /**
     * 获取名称
     */
    public String getName() {
        return name;
    }
    
    /**
     * 重置熔断器
     */
    public void reset() {
        failureCount.set(0);
        lastFailureTime.set(0);
        state.set(State.CLOSED);
        logger.info("Circuit breaker {} has been reset", name);
    }
    
    /**
     * 熔断器状态
     */
    public enum State {
        CLOSED,    // 关闭状态，正常执行
        OPEN,      // 开启状态，拒绝执行
        HALF_OPEN  // 半开状态，尝试执行
    }
    
    /**
     * 熔断器操作接口
     */
    @FunctionalInterface
    public interface CircuitBreakerOperation<T> {
        T execute() throws Exception;
    }
    
    /**
     * 熔断器开启异常
     */
    public static class CircuitBreakerOpenException extends RuntimeException {
        public CircuitBreakerOpenException(String message) {
            super(message);
        }
    }
}